{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "zNmLmqrJAXXp"
   },
   "source": [
    "# EECS 498-007/598-005 Assignment 2-2: Two Layer Neural Network\n",
    "\n",
    "Before we start, please put your name and UMID in following format\n",
    "\n",
    ": Firstname LASTNAME, #00000000   //   e.g.) Justin JOHNSON, #12345678"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "tUGCJrp9Aegm"
   },
   "source": [
    "**Your Answer:**   \n",
    "Your NAME, #XXXXXXXX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "eYE9thuXn4zP"
   },
   "source": [
    "# Setup Code\n",
    "Before getting started, we need to run some boilerplate code to set up our environment, same as Assignment 1. You'll need to rerun this setup code each time you start the notebook.\n",
    "\n",
    "First, run this cell load the autoreload extension. This allows us to edit .py source files, and re-import them into the notebook for a seamless editing and debugging experience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "QqEfH2Rpn9J3"
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CdowvtJen-IP"
   },
   "source": [
    "## Google Colab Setup\n",
    "Next we need to run a few commands to set up our environment on Google Colab. If you are running this notebook on a local machine you can skip this section.\n",
    "\n",
    "Run the following cell to mount your Google Drive. Follow the link, sign in to your Google account (the same account you used to store this notebook!) and copy the authorization code into the text box that appears below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "x5yufsaboBqJ",
    "outputId": "ad3cf435-a39b-4ddd-d7a5-7dff6826471c"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "32Fdd6a5oIT5"
   },
   "source": [
    "Now recall the path in your Google Drive where you uploaded this notebook, fill it in below. If everything is working correctly then running the folowing cell should print the filenames from the assignment:\n",
    "\n",
    "```\n",
    "['two_layer_net.ipynb', 'eecs598', 'two_layer_net.py', 'linear_classifier.py', 'linear_classifier.ipynb', 'a2_helpers.py']\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 54
    },
    "colab_type": "code",
    "id": "MsWH3c7VoL_A",
    "outputId": "85c0f683-eb45-485a-9221-67948facbcf3"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# TODO: Fill in the Google Drive path where you uploaded the assignment\n",
    "# Example: If you create a 2020FA folder and put all the files under A2 folder, then '2020FA/A2'\n",
    "# GOOGLE_DRIVE_PATH_AFTER_MYDRIVE = '2020FA/A2'\n",
    "GOOGLE_DRIVE_PATH_AFTER_MYDRIVE = None\n",
    "GOOGLE_DRIVE_PATH = os.path.join('drive', 'My Drive', GOOGLE_DRIVE_PATH_AFTER_MYDRIVE)\n",
    "print(os.listdir(GOOGLE_DRIVE_PATH))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GaCqHOm9oPB3"
   },
   "source": [
    "Once you have successfully mounted your Google Drive and located the path to this assignment, run th following cell to allow us to import from the `.py` files of this assignment. If it works correctly, it should print the message:\n",
    "\n",
    "```\n",
    "Hello from two_layer_net.py!\n",
    "Hello from a2_helpers.py!\n",
    "```\n",
    "\n",
    "as well as the last edit time for the file `two_layer_net.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "oCaNVx6JoWid",
    "outputId": "2133e4c6-8a6e-4ea3-dd97-23ad471ba2b0"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(GOOGLE_DRIVE_PATH)\n",
    "\n",
    "import time, os\n",
    "os.environ[\"TZ\"] = \"US/Eastern\"\n",
    "time.tzset()\n",
    "\n",
    "from two_layer_net import hello_two_layer_net\n",
    "hello_two_layer_net()\n",
    "\n",
    "from a2_helpers import hello_helper\n",
    "hello_helper()\n",
    "\n",
    "two_layer_net_path = os.path.join(GOOGLE_DRIVE_PATH, 'two_layer_net.py')\n",
    "two_layer_net_edit_time = time.ctime(os.path.getmtime(two_layer_net_path))\n",
    "print('two_layer_net.py last edited on %s' % two_layer_net_edit_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KtMy3qeipNK3"
   },
   "source": [
    "## Miscellaneous"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "z6WjZGY8A9CI"
   },
   "source": [
    "Run some setup code for this notebook: Import some useful packages and increase the default figure size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "O3EvIZ0uAOVN"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import statistics\n",
    "import random\n",
    "import time\n",
    "%matplotlib inline\n",
    "\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (10.0, 8.0)\n",
    "plt.rcParams['font.size'] = 16\n",
    "plt.rcParams['image.interpolation'] = 'nearest'\n",
    "plt.rcParams['image.cmap'] = 'gray'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OvUDZWGU3VLV"
   },
   "source": [
    "We will use GPUs to accelerate our computation in this notebook. Run the following to make sure GPUs are enabled:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "RrAX9FOLpr9k",
    "outputId": "f651b585-1ace-442d-e3d9-e083eb55ae09"
   },
   "outputs": [],
   "source": [
    "if torch.cuda.is_available:\n",
    "  print('Good to go!')\n",
    "else:\n",
    "  print('Please set GPU via Edit -> Notebook Settings.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Hbe3wUpVAjma"
   },
   "source": [
    "# Implementing a Neural Network\n",
    "In this exercise we will develop a neural network with fully-connected layers to perform classification, and test it out on the CIFAR-10 dataset.\n",
    "\n",
    "We train the network with a softmax loss function and L2 regularization on the weight matrices. The network uses a ReLU nonlinearity after the first fully connected layer. \n",
    "\n",
    "In other words, the network has the following architecture:\n",
    "\n",
    "  input - fully connected layer - ReLU - fully connected layer - softmax\n",
    "\n",
    "The outputs of the second fully-connected layer are the scores for each class.\n",
    "\n",
    "**Note**: When you implment the regularization over W, **please DO NOT multiply the regularization term by 1/2** (no coefficient). \n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "lJqim3P1qZgv"
   },
   "source": [
    "## Play with a toy data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5T-4Phbd9GvI"
   },
   "source": [
    "The inputs to our network will be a batch of $N$ (`num_inputs`) $D$-dimensional vectors (`input_size`); the hidden layer will have $H$ hidden units (`hidden_size`), and we will predict classification scores for $C$ categories (`num_classes`). This means that the learnable weights and biases of the network will have the following shapes:\n",
    "\n",
    "*   W1: First layer weights; has shape (D, H)\n",
    "*   b1: First layer biases; has shape (H,)\n",
    "*   W2: Second layer weights; has shape (H, C)\n",
    "*   b2: Second layer biases; has shape (C,)\n",
    "\n",
    "We will use `a2_helpers.get_toy_data` function to generate random weights for a small toy model while we implement the model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZLdCF3B-AOVT"
   },
   "source": [
    "### Forward pass: compute scores\n",
    "Like in the Linear Classifiers exercise, we want to write a function that takes as input the model weights and a batch of images and labels, and returns the loss and the gradient of the loss with respect to each model parameter.\n",
    "\n",
    "However rather than attempting to implement the entire function at once, we will take a staged approach and ask you to implement the full forward and backward pass one step at a time.\n",
    "\n",
    "First we will implement the forward pass of the network which uses the weights and biases to compute scores for all inputs in `nn_forward_pass`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "inlH2l-XEtZQ"
   },
   "source": [
    "Compute the scores and compare with the answer. The distance gap should be smaller than 1e-10."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 289
    },
    "colab_type": "code",
    "id": "tZV9_3ZWAOVU",
    "outputId": "7504b688-c002-4676-c064-29adc38f88a4"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "from a2_helpers import get_toy_data\n",
    "from two_layer_net import nn_forward_pass\n",
    "\n",
    "eecs598.reset_seed(0)\n",
    "toy_X, toy_y, params = get_toy_data()\n",
    "\n",
    "# YOUR_TURN: Implement the score computation part of nn_forward_pass\n",
    "scores, _ = nn_forward_pass(params, toy_X)\n",
    "print('Your scores:')\n",
    "print(scores)\n",
    "print(scores.dtype)\n",
    "print()\n",
    "print('correct scores:')\n",
    "correct_scores = torch.tensor([\n",
    "        [ 9.7003e-08, -1.1143e-07, -3.9961e-08],\n",
    "        [-7.4297e-08,  1.1502e-07,  1.5685e-07],\n",
    "        [-2.5860e-07,  2.2765e-07,  3.2453e-07],\n",
    "        [-4.7257e-07,  9.0935e-07,  4.0368e-07],\n",
    "        [-1.8395e-07,  7.9303e-08,  6.0360e-07]], dtype=torch.float32, device=scores.device)\n",
    "print(correct_scores)\n",
    "print()\n",
    "\n",
    "# The difference should be very small. We get < 1e-10\n",
    "scores_diff = (scores - correct_scores).abs().sum().item()\n",
    "print('Difference between your scores and correct scores: %.2e' % scores_diff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "7XNJ3ydEAOVW"
   },
   "source": [
    "### Forward pass: compute loss\n",
    "Now, we implement the first part of `nn_forward_backward` that computes the data and regularization loss.\n",
    "\n",
    "For the data loss, we will use the softmax loss. For the regularization loss we will use L2 regularization on the weight matrices `W1` and `W2`; we will not apply regularization loss to the bias vectors `b1` and `b2`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "C734SdJGE6xh"
   },
   "source": [
    "First, Let's run the following to check your implementation.\n",
    "\n",
    "We compute the loss for the toy data, and compare with the answer computed by our implementation. The difference between the correct and computed loss should be less than `1e-4`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 68
    },
    "colab_type": "code",
    "id": "wgG6w2uKAOVX",
    "outputId": "e198ce0f-1f05-431e-e724-240aeaa09b5e"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "from a2_helpers import get_toy_data\n",
    "from two_layer_net import nn_forward_backward\n",
    "\n",
    "eecs598.reset_seed(0)\n",
    "toy_X, toy_y, params = get_toy_data()\n",
    "\n",
    "# YOUR_TURN: Implement the loss computation part of nn_forward_backward\n",
    "loss, _ = nn_forward_backward(params, toy_X, toy_y, reg=0.05)\n",
    "print('Your loss: ', loss.item())\n",
    "correct_loss = 1.0986121892929077\n",
    "print('Correct loss: ', correct_loss)\n",
    "diff = (correct_loss - loss).item()\n",
    "\n",
    "# should be very small, we get < 1e-4\n",
    "print('Difference: %.4e' % diff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "vExP-7n3AOVa"
   },
   "source": [
    "### Backward pass\n",
    "Now implement the backward pass for the entire network in `nn_forward_backward`.\n",
    "\n",
    "After doing so, we will use numeric gradient checking to see whether the analytic gradient computed by our backward pass mateches a numeric gradient.\n",
    "\n",
    "We will use the functions `eecs598.grad.compute_numeric_gradient` and `eecs598.grad.rel_error` to help with numeric gradient checking. We can learn more about these functions using the `help` command:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 612
    },
    "colab_type": "code",
    "id": "CJitZg6cS8Sf",
    "outputId": "74491629-ea6e-41a5-dced-bc3b60abd40a"
   },
   "outputs": [],
   "source": [
    "help(eecs598.grad.compute_numeric_gradient)\n",
    "print('-' * 80)\n",
    "help(eecs598.grad.rel_error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "93oOdibtW_Kl"
   },
   "source": [
    "Now we will compute the gradient of the loss with respect to the variables `W1`, `b1`, `W2`, and `b2`. Now that you (hopefully!) have a correctly implemented forward pass, you can debug your backward pass using a numeric gradient check.\n",
    "\n",
    "You should see relative errors less than `1e-4` for all parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 85
    },
    "colab_type": "code",
    "id": "qCEkprvoAOVb",
    "outputId": "3e02d110-e672-4e33-80b8-5d898cfb30ef"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "from a2_helpers import get_toy_data\n",
    "from two_layer_net import nn_forward_backward\n",
    "\n",
    "eecs598.reset_seed(0)\n",
    "\n",
    "reg = 0.05\n",
    "toy_X, toy_y, params = get_toy_data(dtype=torch.float64)\n",
    "\n",
    "# YOUR_TURN: Implement the gradient computation part of nn_forward_backward\n",
    "#            When you implement the gradient computation part, you may need to \n",
    "#            implement the `hidden` output in nn_forward_pass, as well.\n",
    "loss, grads = nn_forward_backward(params, toy_X, toy_y, reg=reg)\n",
    "\n",
    "for param_name, grad in grads.items():\n",
    "  param = params[param_name]\n",
    "  f = lambda w: nn_forward_backward(params, toy_X, toy_y, reg=reg)[0]\n",
    "  grad_numeric = eecs598.grad.compute_numeric_gradient(f, param)\n",
    "  error = eecs598.grad.rel_error(grad, grad_numeric)\n",
    "  print('%s max relative error: %e' % (param_name, error))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "LjAUalCBAOVd"
   },
   "source": [
    "### Train the network\n",
    "To train the network we will use stochastic gradient descent (SGD), similar to the SVM and Softmax classifiers. \n",
    "\n",
    "Look at the function `nn_train` and fill in the missing sections to implement the training procedure. This should be very similar to the training procedure you used for the SVM and Softmax classifiers. \n",
    "\n",
    "You will also have to implement `nn_predict`, as the training process periodically performs prediction to keep track of accuracy over time while the network trains. \n",
    "\n",
    "Once you have implemented the method, run the code below to train a two-layer network on toy data. Your final training loss should be less than 1.0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 545
    },
    "colab_type": "code",
    "id": "Wgw06cLXAOVd",
    "outputId": "be163c99-6590-4354-eafa-93d623bed3a8"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "from a2_helpers import get_toy_data\n",
    "from two_layer_net import nn_forward_backward, nn_train, nn_predict\n",
    "\n",
    "eecs598.reset_seed(0)\n",
    "toy_X, toy_y, params = get_toy_data()\n",
    "\n",
    "# YOUR_TURN: Implement the nn_train function.\n",
    "#            You may need to check nn_predict function (the \"pred_func\") as well.\n",
    "stats = nn_train(params, nn_forward_backward, nn_predict, toy_X, toy_y, toy_X, toy_y,\n",
    "                 learning_rate=1e-1, reg=1e-6,\n",
    "                 num_iters=200, verbose=False)\n",
    "\n",
    "print('Final training loss: ', stats['loss_history'][-1])\n",
    "\n",
    "# plot the loss history\n",
    "plt.plot(stats['loss_history'], 'o')\n",
    "plt.xlabel('Iteration')\n",
    "plt.ylabel('training loss')\n",
    "plt.title('Training Loss history')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 528
    },
    "colab_type": "code",
    "id": "EUS4aDp_HzG1",
    "outputId": "e0c63f1f-2077-499f-c9c4-8a83e09635ae"
   },
   "outputs": [],
   "source": [
    "# Plot the loss function and train / validation accuracies\n",
    "plt.plot(stats['train_acc_history'], 'o', label='train')\n",
    "plt.plot(stats['val_acc_history'], 'o', label='val')\n",
    "plt.title('Classification accuracy history')\n",
    "plt.xlabel('Epoch')\n",
    "plt.ylabel('Clasification accuracy')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8cPIajWNAOVg"
   },
   "source": [
    "## Testing our NN on a real dataset: CIFAR-10\n",
    "Now that you have implemented a two-layer network that passes gradient checks and works on toy data, it's time to load up our favorite CIFAR-10 data so we can use it to train a classifier on a real dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 568
    },
    "colab_type": "code",
    "id": "lYo_XrU3AOVg",
    "outputId": "e0e8ca93-3570-45f4-96ec-d03d91b1148c"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "\n",
    "# Invoke the above function to get our data.\n",
    "eecs598.reset_seed(0)\n",
    "data_dict = eecs598.data.preprocess_cifar10(dtype=torch.float64)\n",
    "print('Train data shape: ', data_dict['X_train'].shape)\n",
    "print('Train labels shape: ', data_dict['y_train'].shape)\n",
    "print('Validation data shape: ', data_dict['X_val'].shape)\n",
    "print('Validation labels shape: ', data_dict['y_val'].shape)\n",
    "print('Test data shape: ', data_dict['X_test'].shape)\n",
    "print('Test labels shape: ', data_dict['y_test'].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Cq-HkgRBAOVQ"
   },
   "source": [
    "### Wrap all function into a Class\n",
    "We will use the class `TwoLayerNet` to represent instances of our network. The network parameters are stored in the instance variable `self.params` where keys are string parameter names and values are PyTorch tensors.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "_CsYAv3uAOVi"
   },
   "source": [
    "### Train a network\n",
    "To train our network we will use SGD. In addition, we will adjust the learning rate with an exponential learning rate schedule as optimization proceeds; after each epoch, we will reduce the learning rate by multiplying it by a decay rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 119
    },
    "colab_type": "code",
    "id": "hgg0QV9DAOVj",
    "outputId": "ac949f3a-edf9-4a54-c47a-4c348c6c89e2"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "from two_layer_net import TwoLayerNet\n",
    "\n",
    "input_size = 3 * 32 * 32\n",
    "hidden_size = 36\n",
    "num_classes = 10\n",
    "\n",
    "# fix random seed before we generate a set of parameters\n",
    "eecs598.reset_seed(0)\n",
    "net = TwoLayerNet(input_size, hidden_size, num_classes, dtype=data_dict['X_train'].dtype, device=data_dict['X_train'].device)\n",
    "\n",
    "# Train the network\n",
    "stats = net.train(data_dict['X_train'], data_dict['y_train'],\n",
    "                  data_dict['X_val'], data_dict['y_val'],\n",
    "                  num_iters=500, batch_size=1000,\n",
    "                  learning_rate=1e-2, learning_rate_decay=0.95,\n",
    "                  reg=0.25, verbose=True)\n",
    "\n",
    "# Predict on the validation set\n",
    "y_val_pred = net.predict(data_dict['X_val'])\n",
    "val_acc = 100.0 * (y_val_pred == data_dict['y_val']).double().mean().item()\n",
    "print('Validation accuracy: %.2f%%' % val_acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ixxgq5RKAOVl"
   },
   "source": [
    "### Debug the training\n",
    "With the default parameters we provided above, you should get a validation accuracy less than 10% on the validation set. This isn't very good.\n",
    "\n",
    "One strategy for getting insight into what's wrong is to plot the loss function and the accuracies on the training and validation sets during optimization.\n",
    "\n",
    "Another strategy is to visualize the weights that were learned in the first layer of the network. In most neural networks trained on visual data, the first layer weights typically show some visible structure when visualized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 310
    },
    "colab_type": "code",
    "id": "6sYXImDTAOVm",
    "outputId": "54a06634-bf14-4bac-e522-b96800b6ae4b"
   },
   "outputs": [],
   "source": [
    "# Plot the loss function and train / validation accuracies\n",
    "from a2_helpers import plot_stats\n",
    "\n",
    "plot_stats(stats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "616EK5UoKgmE"
   },
   "source": [
    "Similar to SVM and Softmax classifier, let's visualize the weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 466
    },
    "colab_type": "code",
    "id": "FnuRjtyKAOVo",
    "outputId": "baeeddfd-8f05-40a9-ddfb-bbbf8736852b"
   },
   "outputs": [],
   "source": [
    "from a2_helpers import show_net_weights \n",
    "\n",
    "show_net_weights(net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OlVbXxmPNzPY"
   },
   "source": [
    "### What's wrong?\n",
    "Looking at the visualizations above, we see that the loss is decreasing more or less linearly, which seems to suggest that the learning rate may be too low. Moreover, there is no gap between the training and validation accuracy, suggesting that the model we used has low capacity, and that we should increase its size. On the other hand, with a very large model we would expect to see more overfitting, which would manifest itself as a very large gap between the training and validation accuracy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rDNZ8ZAnN7hj"
   },
   "source": [
    "#### Capacity?\n",
    "Our initial model has very similar performance on the training and validation sets. This suggests that the model is underfitting, and that its performance might improve if we were to increase its capacity.\n",
    "\n",
    "One way we can increase the capacity of a neural network model is to increase the size of its hidden layer. Here we investigate the effect of increasing the size of the hidden layer. The performance (as measured by validation-set accuracy) should increase as the size of the hidden layer increases; however it may show diminishing returns for larger layer sizes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 450
    },
    "colab_type": "code",
    "id": "_C-ChHUlN68f",
    "outputId": "c5fd64d9-1416-473d-b57d-877fc13f7cd6"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "from a2_helpers import plot_acc_curves\n",
    "from two_layer_net import TwoLayerNet\n",
    "\n",
    "hidden_sizes = [2, 8, 32, 128] \n",
    "lr = 0.1\n",
    "reg = 0.001\n",
    "\n",
    "stat_dict = {}\n",
    "for hs in hidden_sizes:\n",
    "  print('train with hidden size: {}'.format(hs))\n",
    "  # fix random seed before we generate a set of parameters\n",
    "  eecs598.reset_seed(0)\n",
    "  net = TwoLayerNet(3 * 32 * 32, hs, 10, device=data_dict['X_train'].device, dtype=data_dict['X_train'].dtype)\n",
    "  stats = net.train(data_dict['X_train'], data_dict['y_train'], data_dict['X_val'], data_dict['y_val'],\n",
    "            num_iters=3000, batch_size=1000,\n",
    "            learning_rate=lr, learning_rate_decay=0.95,\n",
    "            reg=reg, verbose=False)\n",
    "  stat_dict[hs] = stats\n",
    "\n",
    "plot_acc_curves(stat_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QpSrK3olUfOZ"
   },
   "source": [
    "#### Regularization?\n",
    "Another possible explanation for the small gap we saw between the train and validation accuracies of our model is regularization. In particular, if the regularization coefficient were too high then the model may be unable to fit the training data.\n",
    "\n",
    "We can investigate the phenomenon empirically by training a set of models with varying regularization strengths while fixing other hyperparameters.\n",
    "\n",
    "You should see that setting the regularization strength too high will harm the validation-set performance of the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 433
    },
    "colab_type": "code",
    "id": "DRPsxxFnU3Un",
    "outputId": "2c631a91-a69a-4b62-e36b-29c829536b37"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "from a2_helpers import plot_acc_curves\n",
    "from two_layer_net import TwoLayerNet\n",
    "\n",
    "hs = 128\n",
    "lr = 1.0\n",
    "regs = [0, 1e-5, 1e-3, 1e-1]\n",
    "\n",
    "stat_dict = {}\n",
    "for reg in regs:\n",
    "  print('train with regularization: {}'.format(reg))\n",
    "  # fix random seed before we generate a set of parameters\n",
    "  eecs598.reset_seed(0)\n",
    "  net = TwoLayerNet(3 * 32 * 32, hs, 10, device=data_dict['X_train'].device, dtype=data_dict['X_train'].dtype)\n",
    "  stats = net.train(data_dict['X_train'], data_dict['y_train'], data_dict['X_val'], data_dict['y_val'],\n",
    "            num_iters=3000, batch_size=1000,\n",
    "            learning_rate=lr, learning_rate_decay=0.95,\n",
    "            reg=reg, verbose=False)\n",
    "  stat_dict[reg] = stats\n",
    "\n",
    "plot_acc_curves(stat_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3zFWkxebWXtu"
   },
   "source": [
    "#### Learning Rate?\n",
    "Last but not least, we also want to see the effect of learning rate with respect to the performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 433
    },
    "colab_type": "code",
    "id": "lc_YYCDmWld-",
    "outputId": "19367e7d-0c8e-47fb-e895-a6e6b38934e6"
   },
   "outputs": [],
   "source": [
    "import eecs598\n",
    "from a2_helpers import plot_acc_curves\n",
    "from two_layer_net import TwoLayerNet\n",
    "\n",
    "hs = 128\n",
    "lrs = [1e-4, 1e-2, 1e0, 1e2]\n",
    "reg = 1e-4\n",
    "\n",
    "stat_dict = {}\n",
    "for lr in lrs:\n",
    "  print('train with learning rate: {}'.format(lr))\n",
    "  # fix random seed before we generate a set of parameters\n",
    "  eecs598.reset_seed(0)\n",
    "  net = TwoLayerNet(3 * 32 * 32, hs, 10, device=data_dict['X_train'].device, dtype=data_dict['X_train'].dtype)\n",
    "  stats = net.train(data_dict['X_train'], data_dict['y_train'], data_dict['X_val'], data_dict['y_val'],\n",
    "            num_iters=3000, batch_size=1000,\n",
    "            learning_rate=lr, learning_rate_decay=0.95,\n",
    "            reg=reg, verbose=False)\n",
    "  stat_dict[lr] = stats\n",
    "\n",
    "plot_acc_curves(stat_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mVCEro4FAOVq"
   },
   "source": [
    "### Tune your hyperparameters\n",
    "\n",
    "**Tuning**. Tuning the hyperparameters and developing intuition for how they affect the final performance is a large part of using Neural Networks, so we want you to get a lot of practice. Below, you should experiment with different values of the various hyperparameters, including hidden layer size, learning rate, number of training epochs, and regularization strength. You might also consider tuning the learning rate decay, but you should be able to get good performance using the default value.\n",
    "\n",
    "**Plots**. To guide your hyperparameter search, you might consider making auxiliary plots of training and validation performance as above, or plotting the results arising from different hyperparameter combinations as we did in the Linear Classifier notebook. You should feel free to plot any auxiliary results you need in order to find a good network, but we don't require any particular plots from you.\n",
    "\n",
    "**Approximate results**. To get full credit for the assignment, you should achieve a classification accuracy above 50% on the validation set.\n",
    "\n",
    "(Our best model gets a validation-set accuracy 56.44% -- did you beat us?)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 827
    },
    "colab_type": "code",
    "id": "bG4DjBMIAOVq",
    "outputId": "717919b5-b6da-44dd-c5c9-0b3f03b6b94b"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import eecs598\n",
    "from two_layer_net import TwoLayerNet, find_best_net, nn_get_search_params\n",
    "\n",
    "# running this model on float64 may needs more time, so set it as float32\n",
    "eecs598.reset_seed(0)\n",
    "data_dict = eecs598.data.preprocess_cifar10(dtype=torch.float32)\n",
    "\n",
    "# store the best model into this \n",
    "eecs598.reset_seed(0)\n",
    "best_net, best_stat, best_val_acc = find_best_net(data_dict, nn_get_search_params)\n",
    "print(best_val_acc)\n",
    "\n",
    "plot_stats(best_stat)\n",
    "\n",
    "# save the best model\n",
    "path = os.path.join(GOOGLE_DRIVE_PATH, 'nn_best_model.pt')\n",
    "best_net.save(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "NsYIu49plJ9r",
    "outputId": "a474ddcf-51fd-4e27-fced-38fd6ab0e37d"
   },
   "outputs": [],
   "source": [
    "# Check the validation-set accuracy of your best model\n",
    "y_val_preds = best_net.predict(data_dict['X_val'])\n",
    "val_acc = 100 * (y_val_preds == data_dict['y_val']).double().mean().item()\n",
    "print('Best val-set accuracy: %.2f%%' % val_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 466
    },
    "colab_type": "code",
    "id": "hZgDq4zlAOVt",
    "outputId": "ba3ff655-7c19-4a14-d400-f5bc6094f309"
   },
   "outputs": [],
   "source": [
    "from a2_helpers import show_net_weights\n",
    "# visualize the weights of the best network\n",
    "show_net_weights(best_net)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "UG56gKWsAOVv"
   },
   "source": [
    "### Run on the test set\n",
    "When you are done experimenting, you should evaluate your final trained network on the test set. To get full credit for the assignment, you should achieve over 50% classification accuracy on the test set.\n",
    "\n",
    "(Our best model gets 56.03% test-set accuracy -- did you beat us?)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "2b3h8f8_AOVw",
    "outputId": "35cf6243-78a3-4c5a-b22c-5c9eae0d309d"
   },
   "outputs": [],
   "source": [
    "y_test_preds = best_net.predict(data_dict['X_test'])\n",
    "test_acc = 100 * (y_test_preds == data_dict['y_test']).double().mean().item()\n",
    "print('Test accuracy: %.2f%%' % test_acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "37R_J2uMP3d-"
   },
   "source": [
    "# Submit Your Work\n",
    "After completing both notebooks for this assignment (`linear_classifier.ipynb` and this notebook, `two_layer_net.ipynb`), run the following cell to create a `.zip` file for you to download and turn in. **Please MANUALLY SAVE every `*.ipynb` and `*.py` files before executing the following cell:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "GE2M6t1hP-Pe",
    "outputId": "e7177738-a8a7-43a5-89e2-3258e14ce1e1"
   },
   "outputs": [],
   "source": [
    "from eecs598.submit import make_a2_submission\n",
    "\n",
    "# TODO: Replace these with your actual uniquename and umid\n",
    "uniquename = None\n",
    "umid = None\n",
    "make_a2_submission(GOOGLE_DRIVE_PATH, uniquename, umid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mGbRHnoEAUVN"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "eYE9thuXn4zP",
    "CdowvtJen-IP",
    "KtMy3qeipNK3",
    "Hbe3wUpVAjma",
    "lJqim3P1qZgv",
    "ZLdCF3B-AOVT",
    "7XNJ3ydEAOVW",
    "vExP-7n3AOVa",
    "LjAUalCBAOVd",
    "8cPIajWNAOVg",
    "_CsYAv3uAOVi",
    "ixxgq5RKAOVl",
    "OlVbXxmPNzPY",
    "rDNZ8ZAnN7hj",
    "QpSrK3olUfOZ",
    "3zFWkxebWXtu",
    "mVCEro4FAOVq",
    "UG56gKWsAOVv",
    "37R_J2uMP3d-"
   ],
   "name": "two_layer_net.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
